package cn.myeasyai.rnn.model;

import lombok.extern.slf4j.Slf4j;
import org.datavec.api.io.labels.ParentPathLabelGenerator;
import org.datavec.api.split.FileSplit;
import org.datavec.image.loader.NativeImageLoader;
import org.datavec.image.recordreader.ImageRecordReader;
import org.deeplearning4j.core.storage.StatsStorage;
import org.deeplearning4j.datasets.datavec.RecordReaderDataSetIterator;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.ConvolutionLayer;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.conf.layers.SubsamplingLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.evaluation.classification.Evaluation;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;
import org.nd4j.linalg.learning.config.Nesterovs;
import org.nd4j.linalg.lossfunctions.LossFunctions;
import org.nd4j.linalg.schedule.MapSchedule;
import org.nd4j.linalg.schedule.ScheduleType;

import java.io.File;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;


/**
 * LeNetMNISTReLu 类实现了基于LeNet架构的卷积神经网络，用于MNIST手写数字数据集的分类任务
 * LeNet是经典的卷积神经网络架构，由Yann LeCun提出，适合处理手写数字等小型图像分类问题
 * 本实现使用ReLu激活函数和梯度下降优化器，通过调整学习率调度提高训练效果
 */
@Slf4j // 引入日志注解，方便输出训练过程中的关键信息
public class LeNetMNISTReLu {

    // MNIST数据集下载地址：https://raw.githubusercontent.com/zq2599/blog_download_files/master/files/mnist_png.tar.gz
    // 数据集存放路径（根据实际环境修改）
    public static final String BASE_PATH = "/Users/yangjie/Downloads"; // 本地下载目录（示例）

    /**
     * 主方法：执行MNIST数据集的加载、模型构建、训练、评估和保存
     * @param args 命令行参数（未使用）
     * @throws Exception 处理文件操作和模型训练中可能出现的异常
     */
    public static void main(String[] args) throws Exception {
        // 图像参数配置
        int height = 28; // 图像高度（MNIST数据集图像为28x28像素）
        int width = 28;  // 图像宽度
        int channels = 1; // 颜色通道数（1表示灰度图，RGB图为3）

        // 分类参数配置
        int outputNum = 10; // 输出类别数（MNIST为0-9共10个数字）

        // 训练参数配置
        int batchSize = 54; // 批处理大小：每次训练输入54个样本
        int nEpochs = 1;    // 训练轮次：整个数据集迭代训练1次（实际应用中通常设为10-20）
        int seed = 1234;    // 随机数种子：保证实验可复现（相同种子下随机操作结果一致）

        // 初始化随机数生成器（用于数据集打乱等操作）
        Random randNumGen = new Random(seed);

        // 检查数据集是否存在
        log.info("检查数据集文件夹是否存在：{}", BASE_PATH + "/mnist_png");
        if (!new File(BASE_PATH + "/mnist_png").exists()) {
            log.info("数据集文件不存在，请下载压缩包并解压到：{}", BASE_PATH);
            return; // 数据集不存在时终止程序
        }

        // ====================== 数据预处理 ======================
        // 标签生成器：根据图像文件的父目录名称生成标签（MNIST数据集按数字分类存放文件夹）
        ParentPathLabelGenerator labelMaker = new ParentPathLabelGenerator();

        // 图像归一化处理器：将像素值从[0,255]缩放至[0,1]（神经网络通常需要输入标准化数据）
        DataNormalization imageScaler = new ImagePreProcessingScaler();

        // ---------------------- 训练集处理 ----------------------
        log.info("训练集的矢量化操作...");
        File trainData = new File(BASE_PATH + "/mnist_png/training"); // 训练集文件夹路径
        // 文件分割器：按指定格式（图片格式）分割文件夹中的文件，并用随机数打乱
        FileSplit trainSplit = new FileSplit(trainData, NativeImageLoader.ALLOWED_FORMATS, randNumGen);
        // 图像记录读取器：将图像文件转换为神经网络可处理的记录（高度、宽度、通道数、标签生成器）
        ImageRecordReader trainRR = new ImageRecordReader(height, width, channels, labelMaker);
        trainRR.initialize(trainSplit); // 初始化读取器，加载训练集文件

        // 数据集迭代器：将记录读取器转换为批量迭代器（批大小、标签列索引、总类别数）
        DataSetIterator trainIter = new RecordReaderDataSetIterator(trainRR, batchSize, 1, outputNum);
        imageScaler.fit(trainIter); // 计算训练集的归一化参数（此处实际无操作，仅初始化）
        trainIter.setPreProcessor(imageScaler); // 为训练集应用归一化

        // ---------------------- 测试集处理 ----------------------
        log.info("测试集的矢量化操作...");
        File testData = new File(BASE_PATH + "/mnist_png/testing"); // 测试集文件夹路径
        FileSplit testSplit = new FileSplit(testData, NativeImageLoader.ALLOWED_FORMATS, randNumGen);
        ImageRecordReader testRR = new ImageRecordReader(height, width, channels, labelMaker);
        testRR.initialize(testSplit);
        DataSetIterator testIter = new RecordReaderDataSetIterator(testRR, batchSize, 1, outputNum);
        testIter.setPreProcessor(imageScaler); // 测试集使用与训练集相同的归一化参数

        // ====================== 神经网络配置 ======================
        log.info("配置神经网络");

        // 学习率调度表：随着迭代次数增加，逐步降低学习率（提高收敛精度）
        // 键：迭代次数，值：对应学习率
        Map<Integer, Double> learningRateSchedule = new HashMap<>();
        learningRateSchedule.put(0, 0.06);    // 初始学习率
        learningRateSchedule.put(200, 0.05);  // 200次迭代后调整
        learningRateSchedule.put(600, 0.028); // 600次迭代后调整
        learningRateSchedule.put(800, 0.0060);
        learningRateSchedule.put(1000, 0.001);

        // 构建多层网络配置
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(seed) // 随机种子：保证权重初始化等操作可复现
                .l2(0.0005) // L2正则化系数：防止过拟合（值越大正则化越强）
                // 优化器：使用Nesterov动量梯度下降（加速收敛并减少震荡）
                .updater(new Nesterovs(new MapSchedule(ScheduleType.ITERATION, learningRateSchedule)))
                .weightInit(WeightInit.XAVIER) // 权重初始化方式：Xavier初始化（适合ReLU激活函数）
                .list() // 开始定义网络层

                // 第1层：卷积层（特征提取）
                .layer(new ConvolutionLayer.Builder(5, 5) // 卷积核大小：5x5
                        .nIn(channels) // 输入通道数：1（灰度图）
                        .stride(1, 1) // 步长：1x1（每次滑动1个像素）
                        .nOut(20) // 输出通道数：20个卷积核（提取20种基础特征）
                        .activation(Activation.IDENTITY) // 激活函数：恒等映射（此处暂不激活，后续可替换为ReLU）
                        .build())

                // 第2层：池化层（下采样）
                .layer(new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX) // 最大池化
                        .kernelSize(2, 2) // 池化窗口：2x2
                        .stride(2, 2) // 步长：2x2（输出尺寸变为输入的1/2）
                        .build()) // 作用：减少参数数量，增强平移不变性

                // 第3层：卷积层（高级特征提取）
                .layer(new ConvolutionLayer.Builder(5, 5) // 卷积核大小：5x5
                        .stride(1, 1) // 步长：1x1
                        .nOut(50) // 输出通道数：50（提取更复杂的组合特征）
                        .activation(Activation.IDENTITY) // 激活函数：恒等映射
                        .build())

                // 第4层：池化层（下采样）
                .layer(new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX)
                        .kernelSize(2, 2)
                        .stride(2, 2)
                        .build())

                // 第5层：全连接层（特征整合）
                .layer(new DenseLayer.Builder()
                        .activation(Activation.RELU) // 激活函数：ReLU（引入非线性，解决梯度消失问题）
                        .nOut(500) // 输出神经元数：500（整合前层特征）
                        .build())

                // 第6层：输出层（分类）
                .layer(new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD) // 损失函数：负对数似然（适合多分类）
                        .nOut(outputNum) // 输出神经元数：10（对应10个类别）
                        .activation(Activation.SOFTMAX) // 激活函数：Softmax（输出概率分布）
                        .build())

                .setInputType(InputType.convolutionalFlat(height, width, channels)) // 输入类型：展平的卷积输入（28x28x1）
                .build(); // 完成配置构建

        // 初始化多层神经网络
        MultiLayerNetwork net = new MultiLayerNetwork(conf);
        net.init(); // 初始化网络参数
//初始化用户界面后端

        // 添加训练监听器：每10次迭代输出一次损失值（监控训练进度）
        net.setListeners(new ScoreIterationListener(10));

//        //确保在代码中正确配置了 UI Server
//        UIServer uiServer = UIServer.getInstance();
////        StatsStorage statsStorage = new InMemoryStatsStorage();
//        StatsStorage statsStorage = new FileStatsStorage(new File(BASE_PATH, "ui-stats.dl4j"));
//
//        uiServer.attach(statsStorage);
//        net.setListeners(new StatsListener(statsStorage));
        log.info("神经网络共[{}]个参数", net.numParams()); // 输出网络总参数数量

        // ====================== 模型训练与评估 ======================
        long startTime = System.currentTimeMillis(); // 记录训练开始时间

        // 多轮训练
        for (int i = 0; i < nEpochs; i++) {
            log.info("第[{}]个循环（Epoch）", i);
            net.fit(trainIter); // 用训练集训练模型

            // 每轮训练后用测试集评估
            Evaluation eval = net.evaluate(testIter);
            log.info(eval.stats()); // 输出评估结果（准确率、混淆矩阵等）

            // 重置迭代器（下一轮训练/评估重新遍历数据）
            trainIter.reset();
            testIter.reset();
        }

        log.info("完成训练和测试，耗时[{}]毫秒", System.currentTimeMillis() - startTime);
        final long currentTimeMillis = System.currentTimeMillis();
        // ====================== 模型保存 ======================
        File ministModelPath = new File(BASE_PATH + "/minist-model-"+currentTimeMillis+"+.zip");
        // 保存模型（包括参数和配置，compress=true启用压缩）
        ModelSerializer.writeModel(net, ministModelPath, true);
        log.info("最新的MINIST模型保存在[{}]", ministModelPath.getPath());
        Thread.sleep(Integer.MAX_VALUE);

    }
}